from api_server.db.models.knowledge_base_model import KnowledgeBaseModel, ListKbResponse, KnowledgeBean
from api_server.db.session import with_session, asdict
from configs import logger
from sqlalchemy import desc, asc


@with_session
def add_kb_to_db(session, kb_name, kb_info, vs_type, embed_model, kb_type, create_user_id, create_user_name,
                 kb_nickname=None):
    # 创建知识库实例
    kb = session.query(KnowledgeBaseModel).filter(KnowledgeBaseModel.kb_name.ilike(kb_name)).first()
    if kb_nickname is None:
        kb_nickname = kb_name
    if not kb:
        kb = KnowledgeBaseModel(kb_name=kb_name, kb_info=kb_info, vs_type=vs_type, embed_model=embed_model,
                                kb_type=kb_type, create_user_id=create_user_id, create_user_name=create_user_name,
                                kb_nickname=kb_nickname)
        session.add(kb)
    else:  # update kb with new vs_type and embed_model
        kb.kb_info = kb_info
        kb.vs_type = vs_type
        kb.kb_nickname = kb_nickname
        kb.embed_model = embed_model
    return True


@with_session
def list_kbs_from_db(session, kb_name=None, user_info=None, min_file_count: int = -1):
    result = []
    try:
        # 先读取公共知识库
        public_kb = session.query(KnowledgeBaseModel).filter(KnowledgeBaseModel.kb_type == 'public').order_by(
            desc(KnowledgeBaseModel.kb_prop), desc(KnowledgeBaseModel.id))
        if kb_name:
            public_kb = public_kb.filter(KnowledgeBaseModel.kb_name.ilike(f'%{kb_name}%'))
        public_kb = public_kb.all()
        for kb in public_kb:
            kb = KnowledgeBean.from_orm(kb)
            result.append(kb)
        # 读取跟人知识库
        if user_info is not None:
            kbs = session.query(KnowledgeBaseModel).filter(KnowledgeBaseModel.file_count > min_file_count,
                                                           KnowledgeBaseModel.kb_type == 'person').order_by(
                desc(KnowledgeBaseModel.id))
            if kb_name:
                kbs = kbs.filter(KnowledgeBaseModel.kb_name.ilike(f'%{kb_name}%'))
            if not user_info.is_admin:
                kbs.filter(KnowledgeBaseModel.create_user_id == user_info.user_id)
            kbs = kbs.all()
            for kb in kbs:
                kb = KnowledgeBean.from_orm(kb)
                result.append(kb)
    except Exception as e:
        logger.exception(e)
    return result


@with_session
def kb_exists(session, kb_name):
    kb = session.query(KnowledgeBaseModel).filter(KnowledgeBaseModel.kb_name.ilike(kb_name)).first()
    status = True if kb else False
    return status


@with_session
def kb_exists_person(session, kb_type, user_id):
    '''
    判断该用户是否存在个人知识库
    '''
    kb = session.query(KnowledgeBaseModel).filter(KnowledgeBaseModel.kb_type == kb_type,
                                                  KnowledgeBaseModel.create_user_id == user_id).first()
    status = True if kb else False
    if status:
        return status, kb.kb_name
    return status, None


@with_session
def load_kb_from_db(session, kb_name):
    kb = session.query(KnowledgeBaseModel).filter(KnowledgeBaseModel.kb_name.ilike(kb_name)).first()
    if kb:
        kb_name, vs_type, embed_model = kb.kb_name, kb.vs_type, kb.embed_model
    else:
        kb_name, vs_type, embed_model = None, None, None
    return kb_name, vs_type, embed_model


@with_session
def load_kb_from_id(session, id):
    kb = session.query(KnowledgeBaseModel).filter(KnowledgeBaseModel.id == id).first()
    if kb:
        kb_name, vs_type, embed_model = kb.kb_name, kb.vs_type, kb.embed_model
    else:
        kb_name, vs_type, embed_model = None, None, None
    return kb_name, vs_type, embed_model

@with_session
def get_kb_from_id(session, id):
    kb = session.query(KnowledgeBaseModel).filter(KnowledgeBaseModel.id == id).first()
    if kb:
        kb = KnowledgeBean.from_orm(kb)
        return kb
    return None
@with_session
def get_kb_from_name(session, kb_name):
    kb = session.query(KnowledgeBaseModel).filter(KnowledgeBaseModel.kb_name == kb_name).first()
    if kb:
        kb = KnowledgeBean.from_orm(kb)
        return kb
    return None

@with_session
def delete_kb_from_db(session, kb_name):
    kb = session.query(KnowledgeBaseModel).filter(KnowledgeBaseModel.kb_name.ilike(kb_name)).first()
    if kb:
        session.delete(kb)
    return True


@with_session
def get_kb_detail(session, kb_name: str) -> dict:
    kb: KnowledgeBaseModel = session.query(KnowledgeBaseModel).filter(KnowledgeBaseModel.kb_name.ilike(kb_name)).first()
    if kb:
        return {
            "id": kb.id,
            "kb_name": kb.kb_name,
            "kb_nickname": kb.kb_nickname,
            "kb_info": kb.kb_info,
            "vs_type": kb.vs_type,
            "embed_model": kb.embed_model,
            "file_count": kb.file_count,
            "create_time": kb.create_time,
        }
    else:
        return {}


if __name__ == "__main__":
    temp = list_kbs_from_db()
    print(temp)
